// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/ipc/message.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include <nebula/core/buffer.h>
#include <nebula/core/device.h>
#include <nebula/io/interfaces.h>
#include <nebula/ipc/metadata_internal.h>
#include <nebula/ipc/options.h>
#include <nebula/ipc/reader.h>
#include <nebula/ipc/reader_internal.h>
#include <nebula/ipc/util.h>
#include <turbo/utility/status.h>
#include <nebula/util/endian.h>
#include <nebula/future/future.h>
#include <turbo/log/logging.h>
#include <turbo/base/ubsan.h>

#include <nebula/fbs/Message_generated.h>

namespace nebula {

    class KeyValueMetadata;

    class MemoryPool;

    namespace ipc {

        class Message::MessageImpl {
        public:
            explicit MessageImpl(std::shared_ptr<Buffer> metadata, std::shared_ptr<Buffer> body)
                    : metadata_(std::move(metadata)), message_(nullptr), body_(std::move(body)) {}

            turbo::Status open() {
                TURBO_RETURN_NOT_OK(
                        internal::VerifyMessage(metadata_->data(), metadata_->size(), &message_));

                // Check that the metadata version is supported
                if (message_->version() < internal::kMinMetadataVersion) {
                    return turbo::invalid_argument_error("Old metadata version not supported");
                }

                if (message_->version() > flatbuf::MetadataVersion::MAX) {
                    return turbo::invalid_argument_error("Unsupported future MetadataVersion: ",
                                           static_cast<int16_t>(message_->version()));
                }

                if (message_->custom_metadata() != nullptr) {
                    // deserialize from Flatbuffers if first time called
                    std::shared_ptr<KeyValueMetadata> md;
                    TURBO_RETURN_NOT_OK(internal::GetKeyValueMetadata(message_->custom_metadata(), &md));
                    custom_metadata_ = std::move(md);  // const-ify
                }

                return turbo::OkStatus();
            }

            MessageType type() const {
                switch (message_->header_type()) {
                    case flatbuf::MessageHeader::Schema:
                        return MessageType::SCHEMA;
                    case flatbuf::MessageHeader::DictionaryBatch:
                        return MessageType::DICTIONARY_BATCH;
                    case flatbuf::MessageHeader::RecordBatch:
                        return MessageType::RECORD_BATCH;
                    case flatbuf::MessageHeader::Tensor:
                        return MessageType::TENSOR;
                    case flatbuf::MessageHeader::SparseTensor:
                        return MessageType::SPARSE_TENSOR;
                    default:
                        return MessageType::NONE;
                }
            }

            MetadataVersion version() const {
                return internal::GetMetadataVersion(message_->version());
            }

            const void *header() const { return message_->header(); }

            int64_t body_length() const { return message_->body_length(); }

            std::shared_ptr<Buffer> body() const { return body_; }

            std::shared_ptr<Buffer> metadata() const { return metadata_; }

            const std::shared_ptr<const KeyValueMetadata> &custom_metadata() const {
                return custom_metadata_;
            }

        private:
            // The Flatbuffer metadata
            std::shared_ptr<Buffer> metadata_;
            const flatbuf::Message *message_;

            // The reconstructed custom_metadata field from the Message Flatbuffer
            std::shared_ptr<const KeyValueMetadata> custom_metadata_;

            // The message body, if any
            std::shared_ptr<Buffer> body_;
        };

        Message::Message(std::shared_ptr<Buffer> metadata, std::shared_ptr<Buffer> body) {
            impl_.reset(new MessageImpl(std::move(metadata), std::move(body)));
        }

        turbo::Result<std::unique_ptr<Message>> Message::open(std::shared_ptr<Buffer> metadata,
                                                       std::shared_ptr<Buffer> body) {
            std::unique_ptr<Message> result(new Message(std::move(metadata), std::move(body)));
            TURBO_RETURN_NOT_OK(result->impl_->open());
            return result;
        }

        Message::~Message() {}

        std::shared_ptr<Buffer> Message::body() const { return impl_->body(); }

        int64_t Message::body_length() const { return impl_->body_length(); }

        std::shared_ptr<Buffer> Message::metadata() const { return impl_->metadata(); }

        MessageType Message::type() const { return impl_->type(); }

        MetadataVersion Message::metadata_version() const { return impl_->version(); }

        const void *Message::header() const { return impl_->header(); }

        const std::shared_ptr<const KeyValueMetadata> &Message::custom_metadata() const {
            return impl_->custom_metadata();
        }

        bool Message::equals(const Message &other) const {
            int64_t metadata_bytes = std::min(metadata()->size(), other.metadata()->size());

            if (!metadata()->equals(*other.metadata(), metadata_bytes)) {
                return false;
            }

            // Compare bodies, if they have them
            auto this_body = body();
            auto other_body = other.body();

            const bool this_has_body = (this_body != nullptr) && (this_body->size() > 0);
            const bool other_has_body = (other_body != nullptr) && (other_body->size() > 0);

            if (this_has_body && other_has_body) {
                return this_body->equals(*other_body);
            } else if (this_has_body ^ other_has_body) {
                // One has a body but not the other
                return false;
            } else {
                // Neither has a body
                return true;
            }
        }

        turbo::Status MaybeAlignMetadata(std::shared_ptr<Buffer> *metadata) {
            if (reinterpret_cast<uintptr_t>((*metadata)->data()) % 8 != 0) {
                // If the metadata memory is not aligned, we copy it here to avoid
                // potential UBSAN issues from Flatbuffers
                TURBO_MOVE_OR_RAISE(*metadata, (*metadata)->copy_slice(0, (*metadata)->size()));
            }
            return turbo::OkStatus();
        }

        turbo::Status CheckMetadataAndGetBodyLength(BufferSpan metadata, int64_t *body_length) {
            const flatbuf::Message *fb_message = nullptr;
            TURBO_RETURN_NOT_OK(internal::VerifyMessage(metadata.data(), metadata.size(), &fb_message));
            *body_length = fb_message->body_length();
            if (*body_length < 0) {
                return turbo::io_error("Invalid IPC message: negative body_length");
            }
            return turbo::OkStatus();
        }

        turbo::Result<std::unique_ptr<Message>> Message::read_from(std::shared_ptr<Buffer> metadata,
                                                           io::InputStream *stream) {
            std::unique_ptr<Message> result;
            auto listener = std::make_shared<AssignMessageDecoderListener>(&result);
            MessageDecoder decoder(listener, MessageDecoder::State::METADATA, metadata->size());
            TURBO_RETURN_NOT_OK(decoder.consume(metadata));

            TURBO_MOVE_OR_RAISE(auto body, stream->read(decoder.next_required_size()));
            if (body->size() < decoder.next_required_size()) {
                return turbo::io_error("Expected to be able to read ", decoder.next_required_size(),
                                       " bytes for message body, got ", body->size());
            }
            TURBO_RETURN_NOT_OK(decoder.consume(body));
            return result;
        }

        turbo::Result<std::unique_ptr<Message>> Message::read_from(const int64_t offset,
                                                           std::shared_ptr<Buffer> metadata,
                                                           io::RandomAccessFile *file) {
            std::unique_ptr<Message> result;
            auto listener = std::make_shared<AssignMessageDecoderListener>(&result);
            MessageDecoder decoder(listener, MessageDecoder::State::METADATA, metadata->size());
            TURBO_RETURN_NOT_OK(decoder.consume(metadata));

            TURBO_MOVE_OR_RAISE(auto body, file->read_at(offset, decoder.next_required_size()));
            if (body->size() < decoder.next_required_size()) {
                return turbo::io_error("Expected to be able to read ", decoder.next_required_size(),
                                       " bytes for message body, got ", body->size());
            }
            TURBO_RETURN_NOT_OK(decoder.consume(body));
            return result;
        }

        turbo::Status WritePadding(io::OutputStream *stream, int64_t nbytes) {
            while (nbytes > 0) {
                const int64_t bytes_to_write = std::min<int64_t>(nbytes, kArrowAlignment);
                TURBO_RETURN_NOT_OK(stream->write(kPaddingBytes, bytes_to_write));
                nbytes -= bytes_to_write;
            }
            return turbo::OkStatus();
        }

        turbo::Status Message::serialize_to(io::OutputStream *stream, const IpcWriteOptions &options,
                                    int64_t *output_length) const {
            int32_t metadata_length = 0;
            TURBO_RETURN_NOT_OK(WriteMessage(*metadata(), options, stream, &metadata_length));

            *output_length = metadata_length;

            auto body_buffer = body();
            if (body_buffer) {
                TURBO_RETURN_NOT_OK(stream->write(body_buffer));
                *output_length += body_buffer->size();

                        DKCHECK_GE(this->body_length(), body_buffer->size());

                int64_t remainder = this->body_length() - body_buffer->size();
                TURBO_RETURN_NOT_OK(WritePadding(stream, remainder));
                *output_length += remainder;
            }
            return turbo::OkStatus();
        }

        bool Message::Verify() const {
            const flatbuf::Message *unused;
            return internal::VerifyMessage(metadata()->data(), metadata()->size(), &unused).ok();
        }

        std::string FormatMessageType(MessageType type) {
            switch (type) {
                case MessageType::SCHEMA:
                    return "schema";
                case MessageType::RECORD_BATCH:
                    return "record batch";
                case MessageType::DICTIONARY_BATCH:
                    return "dictionary";
                case MessageType::TENSOR:
                    return "tensor";
                case MessageType::SPARSE_TENSOR:
                    return "sparse tensor";
                default:
                    break;
            }
            return "unknown";
        }

        turbo::Status ReadFieldsSubset(int64_t offset, int32_t metadata_length,
                                io::RandomAccessFile *file,
                                const FieldsLoaderFunction &fields_loader,
                                const std::shared_ptr<Buffer> &metadata, int64_t required_size,
                                std::shared_ptr<Buffer> &body) {
            const flatbuf::Message *message = nullptr;
            uint8_t continuation_metadata_size = sizeof(int32_t) + sizeof(int32_t);
            // skip 8 bytes (32-bit continuation indicator + 32-bit little-endian length prefix)
            TURBO_RETURN_NOT_OK(internal::VerifyMessage(metadata->data() + continuation_metadata_size,
                                                  metadata->size() - continuation_metadata_size,
                                                  &message));
            auto batch = message->header_as_RecordBatch();
            if (batch == nullptr) {
                return turbo::io_error(
                        "Header-type of flatbuffer-encoded Message is not RecordBatch.");
            }
            internal::IoRecordedRandomAccessFile io_recorded_random_access_file(required_size);
            TURBO_RETURN_NOT_OK(fields_loader(batch, &io_recorded_random_access_file));
            auto const &read_ranges = io_recorded_random_access_file.GetReadRanges();
            for (auto const &range: read_ranges) {
                auto read_result = file->read_at(offset + metadata_length + range.offset, range.length,
                                                body->mutable_data() + range.offset);
                if (!read_result.ok()) {
                    return turbo::io_error("Failed to read message body, error ",
                                           read_result.status().to_string());
                }
            }
            return turbo::OkStatus();
        }

        turbo::Result<std::unique_ptr<Message>> ReadMessage(std::shared_ptr<Buffer> metadata,
                                                     std::shared_ptr<Buffer> body) {
            std::unique_ptr<Message> result;
            auto listener = std::make_shared<AssignMessageDecoderListener>(&result);
            // If the user does not pass in a body buffer then we assume they are skipping it
            MessageDecoder decoder(listener, default_memory_pool(), body == nullptr);

            if (metadata->size() < decoder.next_required_size()) {
                return turbo::invalid_argument_error("metadata_length should be at least ",
                                       decoder.next_required_size());
            }

            TURBO_RETURN_NOT_OK(decoder.consume(metadata));

            switch (decoder.state()) {
                case MessageDecoder::State::INITIAL:
                    // Metadata did not request a body so we better not have provided one
                            DKCHECK_EQ(body, nullptr);
                    return result;
                case MessageDecoder::State::METADATA_LENGTH:
                    return turbo::invalid_argument_error("metadata length is missing from the metadata buffer");
                case MessageDecoder::State::METADATA:
                    return turbo::invalid_argument_error("flatbuffer size ", decoder.next_required_size(),
                                           " invalid. Buffer size: ", metadata->size());
                case MessageDecoder::State::BODY: {
                    if (body == nullptr) {
                        // Caller didn't give a body so just give them a message without body
                        return result;
                    }
                    if (body->size() != decoder.next_required_size()) {
                        return turbo::io_error("Expected body buffer to be ",
                                               decoder.next_required_size(),
                                               " bytes for message body, got ", body->size());
                    }
                    TURBO_RETURN_NOT_OK(decoder.consume(body));
                    return result;
                }
                case MessageDecoder::State::EOS:
                    return turbo::invalid_argument_error("Unexpected empty message in IPC file format");
                default:
                    return turbo::invalid_argument_error("Unexpected state: ", decoder.state());
            }
        }

        turbo::Result<std::unique_ptr<Message>> ReadMessage(int64_t offset, int32_t metadata_length,
                                                     io::RandomAccessFile *file,
                                                     const FieldsLoaderFunction &fields_loader) {
            std::unique_ptr<Message> result;
            auto listener = std::make_shared<AssignMessageDecoderListener>(&result);
            MessageDecoder decoder(listener);

            if (metadata_length < decoder.next_required_size()) {
                return turbo::invalid_argument_error("metadata_length should be at least ",
                                       decoder.next_required_size());
            }

            TURBO_MOVE_OR_RAISE(auto metadata, file->read_at(offset, metadata_length));
            if (metadata->size() < metadata_length) {
                return turbo::invalid_argument_error("Expected to read ", metadata_length,
                                       " metadata bytes but got ", metadata->size());
            }
            TURBO_RETURN_NOT_OK(decoder.consume(metadata));

            switch (decoder.state()) {
                case MessageDecoder::State::INITIAL:
                    return result;
                case MessageDecoder::State::METADATA_LENGTH:
                    return turbo::invalid_argument_error("metadata length is missing. File offset: ", offset,
                                           ", metadata length: ", metadata_length);
                case MessageDecoder::State::METADATA:
                    return turbo::invalid_argument_error("flatbuffer size ", decoder.next_required_size(),
                                           " invalid. File offset: ", offset,
                                           ", metadata length: ", metadata_length);
                case MessageDecoder::State::BODY: {
                    std::shared_ptr<Buffer> body;
                    if (fields_loader) {
                        TURBO_MOVE_OR_RAISE(
                                body, allocate_buffer(decoder.next_required_size(), default_memory_pool()));
                        TURBO_RETURN_NOT_OK(ReadFieldsSubset(offset, metadata_length, file, fields_loader,
                                                       metadata, decoder.next_required_size(), body));
                    } else {
                        TURBO_MOVE_OR_RAISE(
                                body, file->read_at(offset + metadata_length, decoder.next_required_size()));
                    }
                    if (body->size() < decoder.next_required_size()) {
                        return turbo::io_error("Expected to be able to read ",
                                               decoder.next_required_size(),
                                               " bytes for message body, got ", body->size());
                    }
                    TURBO_RETURN_NOT_OK(decoder.consume(body));
                    return result;
                }
                case MessageDecoder::State::EOS:
                    return turbo::invalid_argument_error("Unexpected empty message in IPC file format");
                default:
                    return turbo::invalid_argument_error("Unexpected state: ", decoder.state());
            }
        }

        Future<std::shared_ptr<Message>> ReadMessageAsync(int64_t offset, int32_t metadata_length,
                                                          int64_t body_length,
                                                          io::RandomAccessFile *file,
                                                          const io::IOContext &context) {
            struct State {
                std::unique_ptr<Message> result;
                std::shared_ptr<MessageDecoderListener> listener;
                std::shared_ptr<MessageDecoder> decoder;
            };
            auto state = std::make_shared<State>();
            state->listener = std::make_shared<AssignMessageDecoderListener>(&state->result);
            state->decoder = std::make_shared<MessageDecoder>(state->listener);

            if (metadata_length < state->decoder->next_required_size()) {
                return turbo::invalid_argument_error("metadata_length should be at least ",
                                       state->decoder->next_required_size());
            }
            return file->read_async(context, offset, metadata_length + body_length)
                    .Then([=](std::shared_ptr<Buffer> metadata) -> turbo::Result<std::shared_ptr<Message>> {
                        if (metadata->size() < metadata_length) {
                            return turbo::invalid_argument_error("Expected to read ", metadata_length,
                                                   " metadata bytes but got ", metadata->size());
                        }
                        TURBO_RETURN_NOT_OK(
                                state->decoder->consume(SliceBuffer(metadata, 0, metadata_length)));
                        switch (state->decoder->state()) {
                            case MessageDecoder::State::INITIAL:
                                return std::move(state->result);
                            case MessageDecoder::State::METADATA_LENGTH:
                                return turbo::invalid_argument_error("metadata length is missing. File offset: ", offset,
                                                       ", metadata length: ", metadata_length);
                            case MessageDecoder::State::METADATA:
                                return turbo::invalid_argument_error("flatbuffer size ",
                                                       state->decoder->next_required_size(),
                                                       " invalid. File offset: ", offset,
                                                       ", metadata length: ", metadata_length);
                            case MessageDecoder::State::BODY: {
                                auto body = SliceBuffer(metadata, metadata_length, body_length);
                                if (body->size() < state->decoder->next_required_size()) {
                                    return turbo::io_error("Expected to be able to read ",
                                                           state->decoder->next_required_size(),
                                                           " bytes for message body, got ", body->size());
                                }
                                TURBO_RETURN_NOT_OK(state->decoder->consume(body));
                                return std::move(state->result);
                            }
                            case MessageDecoder::State::EOS:
                                return turbo::invalid_argument_error("Unexpected empty message in IPC file format");
                            default:
                                return turbo::invalid_argument_error("Unexpected state: ", state->decoder->state());
                        }
                    });
        }

        turbo::Status AlignStream(io::InputStream *stream, int32_t alignment) {
            TURBO_MOVE_OR_RAISE(int64_t position, stream->tell());
            return stream->advance(PaddedLength(position, alignment) - position);
        }

        turbo::Status AlignStream(io::OutputStream *stream, int32_t alignment) {
            TURBO_MOVE_OR_RAISE(int64_t position, stream->tell());
            int64_t remainder = PaddedLength(position, alignment) - position;
            if (remainder > 0) {
                return stream->write(kPaddingBytes, remainder);
            }
            return turbo::OkStatus();
        }

        turbo::Status CheckAligned(io::FileInterface *stream, int32_t alignment) {
            TURBO_MOVE_OR_RAISE(int64_t position, stream->tell());
            if (position % alignment != 0) {
                return turbo::invalid_argument_error("Stream is not aligned pos: ", position,
                                       " alignment: ", alignment);
            } else {
                return turbo::OkStatus();
            }
        }

        turbo::Status DecodeMessage(MessageDecoder *decoder, io::InputStream *file) {
            if (decoder->state() == MessageDecoder::State::INITIAL) {
                uint8_t continuation[sizeof(int32_t)];
                TURBO_MOVE_OR_RAISE(int64_t bytes_read, file->read(sizeof(int32_t), &continuation));
                if (bytes_read == 0) {
                    // EOS without indication
                    return turbo::OkStatus();
                } else if (bytes_read != decoder->next_required_size()) {
                    return turbo::invalid_argument_error("Corrupted message, only ", bytes_read, " bytes available");
                }
                TURBO_RETURN_NOT_OK(decoder->consume(continuation, bytes_read));
            }

            if (decoder->state() == MessageDecoder::State::METADATA_LENGTH) {
                // Valid IPC message, read the message length now
                uint8_t metadata_length[sizeof(int32_t)];
                TURBO_MOVE_OR_RAISE(int64_t bytes_read,
                                       file->read(sizeof(int32_t), &metadata_length));
                if (bytes_read != decoder->next_required_size()) {
                    return turbo::invalid_argument_error("Corrupted metadata length, only ", bytes_read,
                                           " bytes available");
                }
                TURBO_RETURN_NOT_OK(decoder->consume(metadata_length, bytes_read));
            }

            if (decoder->state() == MessageDecoder::State::EOS) {
                return turbo::OkStatus();
            }

            auto metadata_length = decoder->next_required_size();
            TURBO_MOVE_OR_RAISE(auto metadata, file->read(metadata_length));
            if (metadata->size() != metadata_length) {
                return turbo::invalid_argument_error("Expected to read ", metadata_length, " metadata bytes, but ",
                                       "only read ", metadata->size());
            }
            TURBO_RETURN_NOT_OK(decoder->consume(metadata));

            if (decoder->state() == MessageDecoder::State::BODY) {
                TURBO_MOVE_OR_RAISE(auto body, file->read(decoder->next_required_size()));
                if (body->size() < decoder->next_required_size()) {
                    return turbo::io_error("Expected to be able to read ",
                                           decoder->next_required_size(),
                                           " bytes for message body, got ", body->size());
                }
                TURBO_RETURN_NOT_OK(decoder->consume(body));
            }

            if (decoder->state() == MessageDecoder::State::INITIAL ||
                decoder->state() == MessageDecoder::State::EOS) {
                return turbo::OkStatus();
            } else {
                return turbo::invalid_argument_error("Failed to decode message");
            }
        }

        turbo::Result<std::unique_ptr<Message>> ReadMessage(io::InputStream *file, MemoryPool *pool) {
            std::unique_ptr<Message> message;
            auto listener = std::make_shared<AssignMessageDecoderListener>(&message);
            MessageDecoder decoder(listener, pool);
            TURBO_RETURN_NOT_OK(DecodeMessage(&decoder, file));
            if (!message) {
                return nullptr;
            } else {
                return message;
            }
        }

        turbo::Status WriteMessage(BufferSpan message, const IpcWriteOptions &options,
                            io::OutputStream *file, int32_t *message_length) {
            const int32_t prefix_size = options.write_legacy_ipc_format ? 4 : 8;
            const int32_t flatbuffer_size = static_cast<int32_t>(message.size());

            int32_t padded_message_length = static_cast<int32_t>(
                    PaddedLength(flatbuffer_size + prefix_size, options.alignment));

            int32_t padding = padded_message_length - flatbuffer_size - prefix_size;

            // The returned message size includes the length prefix, the flatbuffer,
            // plus padding
            *message_length = padded_message_length;

            // ARROW-6314: write continuation / padding token
            if (!options.write_legacy_ipc_format) {
                TURBO_RETURN_NOT_OK(file->write(&internal::kIpcContinuationToken, sizeof(int32_t)));
            }

            // write the flatbuffer size prefix including padding in little endian
            int32_t padded_flatbuffer_size =
                    bit_util::ToLittleEndian(padded_message_length - prefix_size);
            TURBO_RETURN_NOT_OK(file->write(&padded_flatbuffer_size, sizeof(int32_t)));

            // write the flatbuffer
            TURBO_RETURN_NOT_OK(file->write(message.data(), flatbuffer_size));
            if (padding > 0) {
                TURBO_RETURN_NOT_OK(file->write(kPaddingBytes, padding));
            }

            return turbo::OkStatus();
        }

        // ----------------------------------------------------------------------
        // Implement MessageDecoder

        turbo::Status MessageDecoderListener::OnInitial() { return turbo::OkStatus(); }

        turbo::Status MessageDecoderListener::OnMetadataLength() { return turbo::OkStatus(); }

        turbo::Status MessageDecoderListener::OnMetadata() { return turbo::OkStatus(); }

        turbo::Status MessageDecoderListener::OnBody() { return turbo::OkStatus(); }

        turbo::Status MessageDecoderListener::on_eos() { return turbo::OkStatus(); }

        static constexpr auto kMessageDecoderNextRequiredSizeInitial = sizeof(int32_t);
        static constexpr auto kMessageDecoderNextRequiredSizeMetadataLength = sizeof(int32_t);

        class MessageDecoder::MessageDecoderImpl {
        public:
            explicit MessageDecoderImpl(std::shared_ptr<MessageDecoderListener> listener,
                                        State initial_state, int64_t initial_next_required_size,
                                        MemoryPool *pool, bool skip_body)
                    : listener_(std::move(listener)),
                      pool_(pool),
                      memory_manager_(CPUDevice::memory_manager(pool_)),
                      state_(initial_state),
                      next_required_size_(initial_next_required_size),
                      chunks_(),
                      buffered_size_(0),
                      metadata_(nullptr),
                      skip_body_(skip_body) {}

            turbo::Status ConsumeData(const uint8_t *data, int64_t size) {
                if (buffered_size_ == 0) {
                    while (size > 0 && size >= next_required_size_) {
                        auto used_size = next_required_size_;
                        switch (state_) {
                            case State::INITIAL:
                                TURBO_RETURN_NOT_OK(ConsumeInitialData(data, next_required_size_));
                                break;
                            case State::METADATA_LENGTH:
                                TURBO_RETURN_NOT_OK(ConsumeMetadataLengthData(data, next_required_size_));
                                break;
                            case State::METADATA: {
                                // We need to copy metadata because it's used in
                                // ConsumeBody(). ConsumeBody() may be called from another
                                // ConsumeData(). We can't assume that the given data for
                                // the current ConsumeData() call is still valid in the
                                // next ConsumeData() call. So we need to copy metadata
                                // here.
                                TURBO_MOVE_OR_RAISE(std::shared_ptr<Buffer> buffer,
                                                       allocate_buffer(next_required_size_, pool_));
                                memcpy(buffer->mutable_data(), data, next_required_size_);
                                TURBO_RETURN_NOT_OK(ConsumeMetadataBuffer(buffer));
                            }
                                break;
                            case State::BODY: {
                                // We don't need to copy the given data for body because
                                // we can assume that a decoded record batch should be
                                // valid only in a listener_->OnMessageDecoded() call. If
                                // the passed message is needed to be valid after the
                                // call, it's a listener_'s responsibility. The listener_
                                // may copy the data for it.
                                auto buffer = std::make_shared<Buffer>(data, next_required_size_);
                                TURBO_RETURN_NOT_OK(ConsumeBodyBuffer(buffer));
                            }
                                break;
                            case State::EOS:
                                return turbo::OkStatus();
                        }
                        data += used_size;
                        size -= used_size;
                    }
                }

                if (size == 0) {
                    return turbo::OkStatus();
                }

                // We need to copy unused data because the given data for the
                // current ConsumeData() call may be invalid in the next
                // ConsumeData() call.
                TURBO_MOVE_OR_RAISE(std::shared_ptr<Buffer> chunk, allocate_buffer(size, pool_));
                memcpy(chunk->mutable_data(), data, size);
                chunks_.push_back(std::move(chunk));
                buffered_size_ += size;
                return ConsumeChunks();
            }

            turbo::Status ConsumeBuffer(std::shared_ptr<Buffer> buffer) {
                if (buffered_size_ == 0) {
                    while (buffer->size() >= next_required_size_) {
                        auto used_size = next_required_size_;
                        switch (state_) {
                            case State::INITIAL:
                                TURBO_RETURN_NOT_OK(ConsumeInitialBuffer(buffer));
                                break;
                            case State::METADATA_LENGTH:
                                TURBO_RETURN_NOT_OK(ConsumeMetadataLengthBuffer(buffer));
                                break;
                            case State::METADATA:
                                if (buffer->size() == next_required_size_) {
                                    return ConsumeMetadataBuffer(buffer);
                                } else {
                                    auto sliced_buffer = SliceBuffer(buffer, 0, next_required_size_);
                                    TURBO_RETURN_NOT_OK(ConsumeMetadataBuffer(sliced_buffer));
                                }
                                break;
                            case State::BODY:
                                if (buffer->size() == next_required_size_) {
                                    return ConsumeBodyBuffer(buffer);
                                } else {
                                    auto sliced_buffer = SliceBuffer(buffer, 0, next_required_size_);
                                    TURBO_RETURN_NOT_OK(ConsumeBodyBuffer(sliced_buffer));
                                }
                                break;
                            case State::EOS:
                                return turbo::OkStatus();
                        }
                        if (buffer->size() == used_size) {
                            return turbo::OkStatus();
                        }
                        buffer = SliceBuffer(buffer, used_size);
                    }
                }

                if (buffer->size() == 0) {
                    return turbo::OkStatus();
                }

                buffered_size_ += buffer->size();
                chunks_.push_back(std::move(buffer));
                return ConsumeChunks();
            }

            int64_t next_required_size() const { return next_required_size_ - buffered_size_; }

            MessageDecoder::State state() const { return state_; }

        private:
            turbo::Status ConsumeChunks() {
                while (state_ != State::EOS) {
                    if (buffered_size_ < next_required_size_) {
                        return turbo::OkStatus();
                    }

                    switch (state_) {
                        case State::INITIAL:
                            TURBO_RETURN_NOT_OK(ConsumeInitialChunks());
                            break;
                        case State::METADATA_LENGTH:
                            TURBO_RETURN_NOT_OK(ConsumeMetadataLengthChunks());
                            break;
                        case State::METADATA:
                            TURBO_RETURN_NOT_OK(ConsumeMetadataChunks());
                            break;
                        case State::BODY:
                            TURBO_RETURN_NOT_OK(ConsumeBodyChunks());
                            break;
                        case State::EOS:
                            return turbo::OkStatus();
                    }
                }

                return turbo::OkStatus();
            }

            turbo::Status ConsumeInitialData(const uint8_t *data, int64_t size) {
                return ConsumeInitial(bit_util::FromLittleEndian(turbo::safe_load_as<int32_t>(data)));
            }

            turbo::Status ConsumeInitialBuffer(const std::shared_ptr<Buffer> &buffer) {
                TURBO_MOVE_OR_RAISE(auto continuation, ConsumeDataBufferInt32(buffer));
                return ConsumeInitial(bit_util::FromLittleEndian(continuation));
            }

            turbo::Status ConsumeInitialChunks() {
                int32_t continuation = 0;
                TURBO_RETURN_NOT_OK(ConsumeDataChunks(sizeof(int32_t), &continuation));
                return ConsumeInitial(bit_util::FromLittleEndian(continuation));
            }

            turbo::Status ConsumeInitial(int32_t continuation) {
                if (continuation == internal::kIpcContinuationToken) {
                    state_ = State::METADATA_LENGTH;
                    next_required_size_ = kMessageDecoderNextRequiredSizeMetadataLength;
                    TURBO_RETURN_NOT_OK(listener_->OnMetadataLength());
                    // Valid IPC message, read the message length now
                    return turbo::OkStatus();
                } else if (continuation == 0) {
                    state_ = State::EOS;
                    next_required_size_ = 0;
                    TURBO_RETURN_NOT_OK(listener_->on_eos());
                    return turbo::OkStatus();
                } else if (continuation > 0) {
                    state_ = State::METADATA;
                    // ARROW-6314: Backwards compatibility for reading old IPC
                    // messages produced prior to version 0.15.0
                    next_required_size_ = continuation;
                    TURBO_RETURN_NOT_OK(listener_->OnMetadata());
                    return turbo::OkStatus();
                } else {
                    return turbo::io_error("Invalid IPC stream: negative continuation token");
                }
            }

            turbo::Status ConsumeMetadataLengthData(const uint8_t *data, int64_t size) {
                return ConsumeMetadataLength(
                        bit_util::FromLittleEndian(turbo::safe_load_as<int32_t>(data)));
            }

            turbo::Status ConsumeMetadataLengthBuffer(const std::shared_ptr<Buffer> &buffer) {
                TURBO_MOVE_OR_RAISE(auto metadata_length, ConsumeDataBufferInt32(buffer));
                return ConsumeMetadataLength(bit_util::FromLittleEndian(metadata_length));
            }

            turbo::Status ConsumeMetadataLengthChunks() {
                int32_t metadata_length = 0;
                TURBO_RETURN_NOT_OK(ConsumeDataChunks(sizeof(int32_t), &metadata_length));
                return ConsumeMetadataLength(bit_util::FromLittleEndian(metadata_length));
            }

            turbo::Status ConsumeMetadataLength(int32_t metadata_length) {
                if (metadata_length == 0) {
                    state_ = State::EOS;
                    next_required_size_ = 0;
                    TURBO_RETURN_NOT_OK(listener_->on_eos());
                    return turbo::OkStatus();
                } else if (metadata_length > 0) {
                    state_ = State::METADATA;
                    next_required_size_ = metadata_length;
                    TURBO_RETURN_NOT_OK(listener_->OnMetadata());
                    return turbo::OkStatus();
                } else {
                    return turbo::io_error("Invalid IPC message: negative metadata length");
                }
            }

            turbo::Status ConsumeMetadataBuffer(const std::shared_ptr<Buffer> &buffer) {
                if (buffer->is_cpu()) {
                    metadata_ = buffer;
                } else {
                    TURBO_MOVE_OR_RAISE(metadata_, Buffer::ViewOrCopy(buffer, memory_manager_));
                }
                return ConsumeMetadata();
            }

            turbo::Status ConsumeMetadataChunks() {
                if (chunks_[0]->size() >= next_required_size_) {
                    if (chunks_[0]->size() == next_required_size_) {
                        if (chunks_[0]->is_cpu()) {
                            metadata_ = std::move(chunks_[0]);
                        } else {
                            TURBO_MOVE_OR_RAISE(metadata_,
                                                   Buffer::ViewOrCopy(chunks_[0], memory_manager_));
                        }
                        chunks_.erase(chunks_.begin());
                    } else {
                        metadata_ = SliceBuffer(chunks_[0], 0, next_required_size_);
                        if (!chunks_[0]->is_cpu()) {
                            TURBO_MOVE_OR_RAISE(metadata_,
                                                   Buffer::ViewOrCopy(metadata_, memory_manager_));
                        }
                        chunks_[0] = SliceBuffer(chunks_[0], next_required_size_);
                    }
                    buffered_size_ -= next_required_size_;
                } else {
                    TURBO_MOVE_OR_RAISE(metadata_, allocate_buffer(next_required_size_, pool_));
                    TURBO_RETURN_NOT_OK(ConsumeDataChunks(next_required_size_, metadata_->mutable_data()));
                }
                return ConsumeMetadata();
            }

            turbo::Status ConsumeMetadata() {
                TURBO_RETURN_NOT_OK(MaybeAlignMetadata(&metadata_));
                int64_t body_length = -1;
                TURBO_RETURN_NOT_OK(CheckMetadataAndGetBodyLength(*metadata_, &body_length));

                state_ = State::BODY;
                next_required_size_ = skip_body_ ? 0 : body_length;
                TURBO_RETURN_NOT_OK(listener_->OnBody());
                if (next_required_size_ == 0) {
                    auto body = std::make_shared<Buffer>(nullptr, 0);
                    return ConsumeBody(&body);
                } else {
                    return turbo::OkStatus();
                }
            }

            turbo::Status ConsumeBodyBuffer(std::shared_ptr<Buffer> buffer) {
                return ConsumeBody(&buffer);
            }

            turbo::Status ConsumeBodyChunks() {
                if (chunks_[0]->size() >= next_required_size_) {
                    auto used_size = next_required_size_;
                    if (chunks_[0]->size() == next_required_size_) {
                        TURBO_RETURN_NOT_OK(ConsumeBody(&chunks_[0]));
                        chunks_.erase(chunks_.begin());
                    } else {
                        auto body = SliceBuffer(chunks_[0], 0, next_required_size_);
                        TURBO_RETURN_NOT_OK(ConsumeBody(&body));
                        chunks_[0] = SliceBuffer(chunks_[0], used_size);
                    }
                    buffered_size_ -= used_size;
                    return turbo::OkStatus();
                } else {
                    TURBO_MOVE_OR_RAISE(std::shared_ptr<Buffer> body,
                                           allocate_buffer(next_required_size_, pool_));
                    TURBO_RETURN_NOT_OK(ConsumeDataChunks(next_required_size_, body->mutable_data()));
                    return ConsumeBody(&body);
                }
            }

            turbo::Status ConsumeBody(std::shared_ptr<Buffer> *buffer) {
                TURBO_MOVE_OR_RAISE(std::unique_ptr<Message> message,
                                       Message::open(metadata_, *buffer));

                TURBO_RETURN_NOT_OK(listener_->OnMessageDecoded(std::move(message)));
                state_ = State::INITIAL;
                next_required_size_ = kMessageDecoderNextRequiredSizeInitial;
                TURBO_RETURN_NOT_OK(listener_->OnInitial());
                return turbo::OkStatus();
            }

            turbo::Result<int32_t> ConsumeDataBufferInt32(const std::shared_ptr<Buffer> &buffer) {
                if (buffer->is_cpu()) {
                    return turbo::safe_load_as<int32_t>(buffer->data());
                } else {
                    TURBO_MOVE_OR_RAISE(auto cpu_buffer, Buffer::ViewOrCopy(buffer, memory_manager_));
                    return turbo::safe_load_as<int32_t>(cpu_buffer->data());
                }
            }

            turbo::Status ConsumeDataChunks(int64_t nbytes, void *out) {
                size_t offset = 0;
                size_t n_used_chunks = 0;
                auto required_size = nbytes;
                std::shared_ptr<Buffer> last_chunk;
                for (auto &chunk: chunks_) {
                    if (!chunk->is_cpu()) {
                        TURBO_MOVE_OR_RAISE(chunk, Buffer::ViewOrCopy(chunk, memory_manager_));
                    }
                    auto data = chunk->data();
                    auto data_size = chunk->size();
                    auto copy_size = std::min(required_size, data_size);
                    memcpy(static_cast<uint8_t *>(out) + offset, data, copy_size);
                    n_used_chunks++;
                    offset += copy_size;
                    required_size -= copy_size;
                    if (required_size == 0) {
                        if (data_size != copy_size) {
                            last_chunk = SliceBuffer(chunk, copy_size);
                        }
                        break;
                    }
                }
                chunks_.erase(chunks_.begin(), chunks_.begin() + n_used_chunks);
                if (last_chunk.get() != nullptr) {
                    chunks_.insert(chunks_.begin(), std::move(last_chunk));
                }
                buffered_size_ -= offset;
                return turbo::OkStatus();
            }

            std::shared_ptr<MessageDecoderListener> listener_;
            MemoryPool *pool_;
            std::shared_ptr<MemoryManager> memory_manager_;
            State state_;
            int64_t next_required_size_;
            std::vector<std::shared_ptr<Buffer>> chunks_;
            int64_t buffered_size_;
            std::shared_ptr<Buffer> metadata_;  // Must be CPU buffer
            bool skip_body_;
        };

        MessageDecoder::MessageDecoder(std::shared_ptr<MessageDecoderListener> listener,
                                       MemoryPool *pool, bool skip_body) {
            impl_.reset(new MessageDecoderImpl(std::move(listener), State::INITIAL,
                                               kMessageDecoderNextRequiredSizeInitial, pool,
                                               skip_body));
        }

        MessageDecoder::MessageDecoder(std::shared_ptr<MessageDecoderListener> listener,
                                       State initial_state, int64_t initial_next_required_size,
                                       MemoryPool *pool, bool skip_body) {
            impl_.reset(new MessageDecoderImpl(std::move(listener), initial_state,
                                               initial_next_required_size, pool, skip_body));
        }

        MessageDecoder::~MessageDecoder() {}

        turbo::Status MessageDecoder::consume(const uint8_t *data, int64_t size) {
            return impl_->ConsumeData(data, size);
        }

        turbo::Status MessageDecoder::consume(std::shared_ptr<Buffer> buffer) {
            return impl_->ConsumeBuffer(buffer);
        }

        int64_t MessageDecoder::next_required_size() const { return impl_->next_required_size(); }

        MessageDecoder::State MessageDecoder::state() const { return impl_->state(); }

        // ----------------------------------------------------------------------
        // Implement InputStream message reader

        /// \brief Implementation of MessageReader that reads from InputStream
        class InputStreamMessageReader : public MessageReader, public MessageDecoderListener {
        public:
            explicit InputStreamMessageReader(io::InputStream *stream)
                    : stream_(stream),
                      owned_stream_(),
                      message_(),
                      decoder_(std::shared_ptr<InputStreamMessageReader>(this, [](void *) {})) {}

            explicit InputStreamMessageReader(const std::shared_ptr<io::InputStream> &owned_stream)
                    : InputStreamMessageReader(owned_stream.get()) {
                owned_stream_ = owned_stream;
            }

            ~InputStreamMessageReader() {}

            turbo::Status OnMessageDecoded(std::unique_ptr<Message> message) override {
                message_ = std::move(message);
                return turbo::OkStatus();
            }

            turbo::Result<std::unique_ptr<Message>> ReadNextMessage() override {
                TURBO_RETURN_NOT_OK(DecodeMessage(&decoder_, stream_));
                return std::move(message_);
            }

        private:
            io::InputStream *stream_;
            std::shared_ptr<io::InputStream> owned_stream_;
            std::unique_ptr<Message> message_;
            MessageDecoder decoder_;
        };

        std::unique_ptr<MessageReader> MessageReader::open(io::InputStream *stream) {
            return std::make_unique<InputStreamMessageReader>(stream);
        }

        std::unique_ptr<MessageReader> MessageReader::open(
                const std::shared_ptr<io::InputStream> &owned_stream) {
            return std::make_unique<InputStreamMessageReader>(owned_stream);
        }

    }  // namespace ipc
}  // namespace nebula
